"""
Visualization Module: PBPK Modeling Results Chart Generation

This module implements visualization functionality for PBPK modeling results, including:
- Plasma concentration-time curves
- Tissue distribution charts
- Absorption kinetics charts
- Elimination kinetics charts
- Population variability analysis charts
- Sensitivity analysis charts

Author: BlackCat@CPPO
Version: 1.0.0
"""

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from matplotlib.patches import Rectangle
import matplotlib.patches as mpatches
from typing import Dict, List, Optional, Tuple, Any
import warnings

class PlotGenerator:
    """
    PBPK modeling results chart generator class
    
    This class generates various charts for PBPK modeling results, supporting multiple chart types and style customization.
    """
    
    def __init__(self, style: str = "seaborn", figsize: Tuple[int, int] = (10, 6)):
        """
        Initialize chart generator
        
        Parameters:
            style (str): Chart style, options: seaborn, ggplot, classic
            figsize (tuple): Default chart size (width, height)
        """
        self.style = style
        self.figsize = figsize
        self.colors = self._get_color_palette()
        self.setup_style()
    
    def setup_style(self):
        """Set chart style"""
        try:
            plt.style.use(self.style)
        except:
            plt.style.use('default')
        
        # Set font support - handle different system fonts
        try:
            # Common fonts for Linux systems
            plt.rcParams['font.sans-serif'] = ['DejaVu Sans', 'Arial', 'Helvetica', 'sans-serif']
            plt.rcParams['axes.unicode_minus'] = False
            
            # Test font
            fig, ax = plt.subplots()
            ax.text(0.5, 0.5, 'Test', fontsize=12)
            plt.close(fig)
        except Exception as e:
            # Use default font if issues occur
            plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
            warnings.warn(f"Font setup failed, using default font: {e}")
        
        sns.set_palette(self.colors)
    
    def _get_color_palette(self) -> List[str]:
        """Get color palette"""
        return [
            '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd',
            '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf',
            '#aec7e8', '#ffbb78', '#98df8a', '#ff9896', '#c5b0d5'
        ]
    
    def plot_plasma_concentration(self, time_points: np.ndarray, concentrations: np.ndarray, 
                                title: str = "Plasma Concentration-Time Curve", save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot plasma concentration-time curve
        
        Parameters:
            time_points: Time points array (h)
            concentrations: Plasma concentration array (mg/L)
            title: Chart title
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        fig, ax = plt.subplots(figsize=self.figsize)
        
        ax.plot(time_points, concentrations, linewidth=2, color=self.colors[0], label='Plasma Concentration')
        ax.set_xlabel('Time (hours)')
        ax.set_ylabel('Plasma Concentration (mg/L)')
        ax.set_title('Plasma Concentration-Time Curve')
        ax.grid(True, alpha=0.3)
        ax.legend()
        
        # Add key parameter annotations
        if len(concentrations) > 0:
            max_conc = np.max(concentrations)
            max_time = time_points[np.argmax(concentrations)]
            ax.scatter(max_time, max_conc, color='red', s=50, zorder=5)
            ax.annotate(f'Cmax: {max_conc:.2f} mg/L\nTmax: {max_time:.1f} h',
                       xy=(max_time, max_conc), xytext=(max_time + 1, max_conc),
                       arrowprops=dict(arrowstyle='->', color='red'))
        
        plt.tight_layout()
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        return fig
    
    def plot_tissue_concentrations(self, time_points: np.ndarray, tissue_data: Dict[str, np.ndarray],
                                 title: str = "Tissue Concentration-Time Curves", save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot multi-tissue concentration-time curves
        
        Parameters:
            time_points: Time points array (h)
            tissue_data: Tissue concentration dictionary {tissue_name: concentrations}
            title: Chart title
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        fig, ax = plt.subplots(figsize=self.figsize)
        
        for i, (tissue, concentrations) in enumerate(tissue_data.items()):
            color_idx = i % len(self.colors)
            ax.plot(time_points, concentrations, linewidth=2, 
                   color=self.colors[color_idx], label=tissue)
        
        ax.set_xlabel('Time (hours)')
        ax.set_ylabel('Tissue Concentration (mg/kg)')
        ax.set_title('Tissue Concentration-Time Curves')
        ax.grid(True, alpha=0.3)
        ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
        
        plt.tight_layout()
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        return fig
    
    def plot_absorption_profile(self, time_points: np.ndarray, absorption_data: Dict[str, np.ndarray],
                              title: str = "Absorption Kinetics Profile", save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot drug absorption kinetics curves
        
        Parameters:
            time_points: Time points array (h)
            absorption_data: Absorption data dictionary
            title: Chart title
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))
        
        # Cumulative absorption
        ax1.plot(time_points, absorption_data.get('cumulative', []), linewidth=2, color=self.colors[0])
        ax1.set_xlabel('Time (hours)')
        ax1.set_ylabel('Cumulative Absorption (mg)')
        ax1.set_title('Cumulative Absorption vs Time')
        ax1.grid(True, alpha=0.3)
        
        # Absorption rate
        ax2.plot(time_points, absorption_data.get('rate', []), linewidth=2, color=self.colors[1])
        ax2.set_xlabel('Time (hours)')
        ax2.set_ylabel('Absorption Rate (mg/h)')
        ax2.set_title('Absorption Rate vs Time')
        ax2.grid(True, alpha=0.3)
        
        # Fraction absorbed
        ax3.plot(time_points, absorption_data.get('fraction_absorbed', []), linewidth=2, color=self.colors[2])
        ax3.set_xlabel('Time (hours)')
        ax3.set_ylabel('Fraction Absorbed')
        ax3.set_title('Fraction Absorbed vs Time')
        ax3.grid(True, alpha=0.3)
        ax3.set_ylim(0, 1.1)
        
        # Gastric emptying and intestinal transit
        gastric = absorption_data.get('gastric_amount', [])
        intestinal = absorption_data.get('intestinal_amount', [])
        if gastric is not None and intestinal is not None:
            ax4.plot(time_points, gastric, linewidth=2, label='Gastric Drug', color=self.colors[3])
            ax4.plot(time_points, intestinal, linewidth=2, label='Intestinal Drug', color=self.colors[4])
            ax4.set_xlabel('Time (hours)')
            ax4.set_ylabel('Drug Amount (mg)')
            ax4.set_title('Gastrointestinal Drug Transit')
            ax4.legend()
            ax4.grid(True, alpha=0.3)
        
        fig.suptitle('Absorption Kinetics Profile', fontsize=16)
        plt.tight_layout()
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        return fig
    
    def plot_elimination_profile(self, time_points: np.ndarray, elimination_data: Dict[str, Any],
                               title: str = "Elimination Kinetics Profile", save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot drug elimination kinetics curves
        
        Parameters:
            time_points: Time points array (h)
            elimination_data: Elimination data dictionary
            title: Chart title
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))
        
        # Plasma concentration (semi-log)
        plasma_conc = elimination_data.get('plasma_concentration', [])
        if len(plasma_conc) > 0:
            ax1.semilogy(time_points, plasma_conc, linewidth=2, color=self.colors[0])
            ax1.set_xlabel('Time (hours)')
            ax1.set_ylabel('Plasma Concentration (mg/L) (log scale)')
            ax1.set_title('Plasma Concentration-Time Curve (Semi-log)')
            ax1.grid(True, alpha=0.3)
        
        # Clearance composition
        clearance_rates = elimination_data.get('clearance_rates', {})
        if clearance_rates is not None and len(clearance_rates) > 0:
            labels = list(clearance_rates.keys())
            values = list(clearance_rates.values())
            ax2.pie(values, labels=labels, autopct='%1.1f%%', colors=self.colors[:len(values)])
            ax2.set_title('Clearance Composition')
        
        # Tissue elimination
        tissue_elimination = elimination_data.get('tissue_elimination', {})
        for i, (tissue, conc) in enumerate(tissue_elimination.items()):
            color_idx = i % len(self.colors)
            ax3.semilogy(time_points, conc, linewidth=2, color=self.colors[color_idx], label=tissue)
        ax3.set_xlabel('Time (hours)')
        ax3.set_ylabel('Tissue Concentration (mg/kg) (log scale)')
        ax3.set_title('Tissue Elimination Curves')
        ax3.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize=8)
        ax3.grid(True, alpha=0.3)
        
        # Metabolite formation
        metabolite_data = elimination_data.get('metabolite_formation', {})
        metabolite_conc = metabolite_data.get('metabolite_concentration', [])
        if len(metabolite_conc) > 0:
            ax4.plot(time_points, metabolite_conc, linewidth=2, color=self.colors[5], label='Metabolite')
            ax4.set_xlabel('Time (hours)')
            ax4.set_ylabel('Metabolite Concentration (mg/L)')
            ax4.set_title('Metabolite Formation Curve')
            ax4.legend()
            ax4.grid(True, alpha=0.3)
        
        fig.suptitle('Elimination Kinetics Profile', fontsize=16)
        plt.tight_layout()
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        return fig
    
    def plot_population_distribution(self, population_data: Dict[str, np.ndarray],
                                   title: str = "Population Distribution Analysis", save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot population distribution charts
        
        Parameters:
            population_data: Population data dictionary
            title: Chart title
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        fig, axes = plt.subplots(2, 2, figsize=(15, 10))
        axes = axes.flatten()
        
        param_names = list(population_data.keys())
        for i, param in enumerate(param_names[:4]):
            data = population_data[param]
            if len(data) > 0:
                axes[i].hist(data, bins=20, alpha=0.7, color=self.colors[i])
                axes[i].axvline(np.mean(data), color='red', linestyle='--', label=f'Mean: {np.mean(data):.2f}')
                axes[i].axvline(np.median(data), color='blue', linestyle='--', label=f'Median: {np.median(data):.2f}')
                axes[i].set_xlabel(param)
                axes[i].set_ylabel('Frequency')
                axes[i].set_title(f'{param} Distribution')
                axes[i].legend()
                axes[i].grid(True, alpha=0.3)
        
        fig.suptitle('Population Distribution Analysis', fontsize=16)
        plt.tight_layout()
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        return fig
    
    def plot_sensitivity_analysis(self, sensitivity_data: Dict[str, Dict[str, float]],
                                title: str = "Sensitivity Analysis Results", save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot sensitivity analysis results
        
        Parameters:
            sensitivity_data: Sensitivity data dictionary
            title: Chart title
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        fig, ax = plt.subplots(figsize=self.figsize)
        
        parameters = list(sensitivity_data.keys())
        auc_sensitivity = [sensitivity_data[p].get('AUC', 0) for p in parameters]
        cmax_sensitivity = [sensitivity_data[p].get('Cmax', 0) for p in parameters]
        
        x = np.arange(len(parameters))
        width = 0.35
        
        bars1 = ax.bar(x - width/2, auc_sensitivity, width, label='AUC Sensitivity', color=self.colors[0])
        bars2 = ax.bar(x + width/2, cmax_sensitivity, width, label='Cmax Sensitivity', color=self.colors[1])
        
        ax.set_xlabel('Parameters')
        ax.set_ylabel('Sensitivity Coefficient')
        ax.set_title('Sensitivity Analysis Results')
        ax.set_xticks(x)
        ax.set_xticklabels(parameters, rotation=45, ha='right')
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        # Add value annotations
        for bars in [bars1, bars2]:
            for bar in bars:
                height = bar.get_height()
                ax.annotate(f'{height:.2f}',
                           xy=(bar.get_x() + bar.get_width() / 2, height),
                           xytext=(0, 3), textcoords="offset points",
                           ha='center', va='bottom', fontsize=8)
        
        plt.tight_layout()
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        return fig
    
    def plot_model_validation(self, predicted: np.ndarray, observed: np.ndarray,
                            title: str = "Model Validation", save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot model validation charts
        
        Parameters:
            predicted: Predicted values array
            observed: Observed values array
            title: Chart title
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # Scatter plot
        ax1.scatter(observed, predicted, alpha=0.6)
        min_val = min(np.min(observed), np.min(predicted))
        max_val = max(np.max(observed), np.max(predicted))
        ax1.plot([min_val, max_val], [min_val, max_val], 'r--', lw=2)
        ax1.set_xlabel('Observed (mg/L)')
        ax1.set_ylabel('Predicted (mg/L)')
        ax1.set_title(f'{title} - Scatter Plot')
        ax1.grid(True, alpha=0.3)
        
        # Residual plot
        residuals = observed - predicted
        ax2.scatter(predicted, residuals, alpha=0.6)
        ax2.axhline(y=0, color='r', linestyle='--')
        ax2.set_xlabel('Predicted (mg/L)')
        ax2.set_ylabel('Residuals (mg/L)')
        ax2.set_title(f'{title} - Residual Plot')
        ax2.grid(True, alpha=0.3)
        
        plt.tight_layout()
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        return fig
    
    def plot_pbpk_summary_dashboard(self, pbpk_results: Dict[str, Any], save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot PBPK modeling results summary dashboard
        
        Parameters:
            pbpk_results: PBPK modeling results dictionary
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        fig = plt.figure(figsize=(20, 15))
        gs = fig.add_gridspec(3, 3, hspace=0.3, wspace=0.3)
        
        time_points = pbpk_results.get('time_points', [])
        plasma_conc = pbpk_results.get('plasma_concentration', [])
        tissue_concs = pbpk_results.get('tissue_concentrations', {})
        absorption_data = pbpk_results.get('absorption', {})
        elimination_data = pbpk_results.get('elimination', {})
        population_data = pbpk_results.get('population_data', {})
        
        # 1. Plasma concentration-time curve
        ax1 = fig.add_subplot(gs[0, :2])
        if len(plasma_conc) > 0 and len(time_points) > 0:
            ax1.plot(time_points, plasma_conc, linewidth=3, color=self.colors[0])
            ax1.set_xlabel('Time (hours)')
            ax1.set_ylabel('Plasma Concentration (mg/L)')
            ax1.set_title('Plasma Concentration-Time Curve')
            ax1.grid(True, alpha=0.3)
        
        # 2. Tissue concentration heatmap
        ax2 = fig.add_subplot(gs[0, 2])
        if tissue_concs is not None and len(tissue_concs) > 0:
            tissues = list(tissue_concs.keys())[:min(8, len(tissue_concs))]
            concentrations = [tissue_concs[t][-1] if len(tissue_concs[t]) > 0 else 0 for t in tissues]
            y_pos = np.arange(len(tissues))
            ax2.barh(y_pos, concentrations, color=self.colors[:len(tissues)])
            ax2.set_yticks(y_pos)
            ax2.set_yticklabels(tissues)
            ax2.set_xlabel('Concentration (mg/kg)')
            ax2.set_title('Final Tissue Concentrations')
        
        # 3. Absorption kinetics
        ax3 = fig.add_subplot(gs[1, 0])
        cumulative_abs = absorption_data.get('cumulative', [])
        if len(cumulative_abs) > 0 and len(time_points) > 0:
            ax3.plot(time_points, cumulative_abs, linewidth=2, color=self.colors[1])
            ax3.set_xlabel('Time (hours)')
            ax3.set_ylabel('Cumulative Absorption (mg)')
            ax3.set_title('Absorption Kinetics')
            ax3.grid(True, alpha=0.3)
        
        # 4. Clearance composition pie chart
        ax4 = fig.add_subplot(gs[1, 1])
        clearance_rates = elimination_data.get('clearance_rates', {})
        if clearance_rates is not None and len(clearance_rates) > 0:
            labels = list(clearance_rates.keys())
            values = list(clearance_rates.values())
            ax4.pie(values, labels=labels, autopct='%1.1f%%', colors=self.colors[:len(values)])
            ax4.set_title('Clearance Composition')
        
        # 5. Population variability box plot
        ax5 = fig.add_subplot(gs[1, 2])
        if population_data:
            param_name = list(population_data.keys())[0] if population_data else "Parameter"
            data = population_data.get(param_name, [])
            if len(data) > 0:
                ax5.boxplot(data)
                ax5.set_ylabel(param_name)
                ax5.set_title(f'{param_name} Population Distribution')
        
        # 6. Tissue concentration-time curves (top 3 tissues)
        ax6 = fig.add_subplot(gs[2, :])
        if tissue_concs is not None and len(tissue_concs) > 0:
            top_tissues = list(tissue_concs.keys())[:min(3, len(tissue_concs))]
            for i, tissue in enumerate(top_tissues):
                conc = tissue_concs[tissue]
                if len(conc) > 0 and len(time_points) > 0:
                    ax6.plot(time_points, conc, linewidth=2, color=self.colors[i], label=tissue)
            ax6.set_xlabel('Time (hours)')
            ax6.set_ylabel('Tissue Concentration (mg/kg)')
            ax6.set_title('Major Tissue Concentration-Time Curves')
            ax6.legend()
            ax6.grid(True, alpha=0.3)
        
        fig.suptitle('PBPK Modeling Summary Dashboard', fontsize=20)
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        return fig
    
    def plot_concentration_time(self, results: Dict[str, Any], save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot concentration-time curve
        
        Parameters:
            results: Simulation results dictionary
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        # Extract all individual plasma concentration data
        time_points = None
        all_concentrations = []
        
        for individual_id, individual_result in results.items():
            if 'concentration' in individual_result:
                concentrations = individual_result['concentration']
                if time_points is None:
                    time_points = individual_result.get('time', np.arange(len(concentrations)))
                all_concentrations.append(concentrations)
        
        if len(all_concentrations) == 0:
            # Create empty chart if no individual data
            fig, ax = plt.subplots(figsize=self.figsize)
            ax.text(0.5, 0.5, 'No Data', ha='center', va='center')
            ax.set_title('Plasma Concentration-Time Curve')
            return fig
        
        # Calculate mean and standard deviation
        concentrations_array = np.array(all_concentrations)
        mean_concentrations = np.mean(concentrations_array, axis=0)
        std_concentrations = np.std(concentrations_array, axis=0)
        
        # Create chart
        fig, ax = plt.subplots(figsize=self.figsize)
        
        # Plot mean curve
        ax.plot(time_points, mean_concentrations, linewidth=2, color=self.colors[0], label='Mean')
        
        # Plot confidence interval
        upper = mean_concentrations + std_concentrations
        lower = np.maximum(mean_concentrations - std_concentrations, 0)
        ax.fill_between(time_points, lower, upper, alpha=0.2, color=self.colors[0], label='±1σ')
        
        # Plot all individual curves
        for i, conc in enumerate(all_concentrations[:10]):  # Limit to first 10 individuals
            ax.plot(time_points, conc, alpha=0.3, color=self.colors[1], linewidth=1)
        
        ax.set_xlabel('Time (hours)')
        ax.set_ylabel('Plasma Concentration (mg/L)')
        ax.set_title('Plasma Concentration-Time Curve')
        ax.grid(True, alpha=0.3)
        ax.legend()
        
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        
        return fig
    
    def plot_adme_processes(self, results: Dict[str, Any], save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot absorption, distribution, metabolism, and elimination processes
        
        Parameters:
            results: Simulation results dictionary
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        # Get data for first individual
        first_individual = next(iter(results.values()))
        time_points = first_individual.get('time', [])
        
        if time_points is None or len(time_points) == 0:
            fig, ax = plt.subplots(figsize=self.figsize)
            ax.text(0.5, 0.5, 'No Data', ha='center', va='center')
            ax.set_title('ADME Process Analysis')
            return fig
        
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))
        
        # 1. Absorption process
        if 'absorption' in first_individual:
            abs_data = first_individual['absorption']
            if 'cumulative' in abs_data:
                ax1.plot(time_points, abs_data['cumulative'], linewidth=2, color=self.colors[0])
                ax1.set_xlabel('Time (hours)')
                ax1.set_ylabel('Cumulative Absorption (mg)')
                ax1.set_title('Absorption Process')
                ax1.grid(True, alpha=0.3)
        
        # 2. Distribution process
        if 'distribution' in first_individual:
            dist_data = first_individual['distribution']
            if 'volume' in dist_data:
                ax2.bar(['Volume of Distribution'], [dist_data['volume']], color=self.colors[1])
                ax2.set_ylabel('Volume of Distribution (L)')
                ax2.set_title('Distribution Characteristics')
        
        # 3. Elimination process
        if 'elimination' in first_individual:
            elim_data = first_individual['elimination']
            if 'clearance_rates' in elim_data:
                rates = elim_data['clearance_rates']
                labels = list(rates.keys())
                values = list(rates.values())
                ax3.pie(values, labels=labels, autopct='%1.1f%%', colors=self.colors[:len(values)])
                ax3.set_title('Clearance Composition')
        
        # 4. Plasma concentration curve
        if 'concentration' in first_individual:
            conc = first_individual['concentration']
            ax4.plot(time_points, conc, linewidth=2, color=self.colors[3])
            ax4.set_xlabel('Time (hours)')
            ax4.set_ylabel('Plasma Concentration (mg/L)')
            ax4.set_title('Plasma Concentration-Time Curve')
            ax4.grid(True, alpha=0.3)
        
        plt.suptitle('ADME Process Analysis', fontsize=16)
        plt.tight_layout()
        
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        
        return fig
    
    def plot_population_variability(self, results: Dict[str, Any], save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot population variability analysis
        
        Parameters:
            results: Simulation results dictionary
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        # Collect all individual plasma concentration data
        all_concentrations = []
        time_points = None
        
        for individual_id, individual_result in results.items():
            if 'concentration' in individual_result:
                concentrations = individual_result['concentration']
                all_concentrations.append(concentrations)
                if time_points is None:
                    time_points = individual_result.get('time', np.arange(len(concentrations)))
        
        if len(all_concentrations) == 0:
            fig, ax = plt.subplots(figsize=self.figsize)
            ax.text(0.5, 0.5, 'No Data', ha='center', va='center')
            ax.set_title('Population Variability Analysis')
            return fig
        
        fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 10))
        
        concentrations_array = np.array(all_concentrations)
        
        # 1. Concentration distribution boxplot
        selected_times = [0, len(time_points)//2, -1]  # Start, middle, end time points
        for i, time_idx in enumerate(selected_times):
            if time_idx < len(time_points):
                conc_at_time = concentrations_array[:, time_idx]
                ax1.boxplot(conc_at_time, positions=[i], labels=[f'T={time_points[time_idx]:.1f}h'])
        ax1.set_ylabel('Plasma Concentration (mg/L)')
        ax1.set_title('Concentration Distribution at Different Time Points')
        ax1.grid(True, alpha=0.3)
        
        # 2. Cmax distribution histogram
        cmax_values = np.max(concentrations_array, axis=1)
        ax2.hist(cmax_values, bins=20, alpha=0.7, color=self.colors[1])
        ax2.axvline(np.mean(cmax_values), color='red', linestyle='--', label=f'Mean: {np.mean(cmax_values):.2f}')
        ax2.set_xlabel('Cmax (mg/L)')
        ax2.set_ylabel('Frequency')
        ax2.set_title('Cmax Population Distribution')
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        
        # 3. AUC distribution
        auc_values = [np.trapezoid(conc, time_points) for conc in all_concentrations]
        ax3.hist(auc_values, bins=20, alpha=0.7, color=self.colors[2])
        ax3.axvline(np.mean(auc_values), color='red', linestyle='--', label=f'Mean: {np.mean(auc_values):.2f}')
        ax3.set_xlabel('AUC (mg·h/L)')
        ax3.set_ylabel('Frequency')
        ax3.set_title('AUC Population Distribution')
        ax3.legend()
        ax3.grid(True, alpha=0.3)
        
        # 4. Coefficient of variation
        mean_curve = np.mean(concentrations_array, axis=0)
        std_curve = np.std(concentrations_array, axis=0)
        # Avoid division by zero and handle near-zero values
        mask = mean_curve > 0.001  # Only calculate CV for meaningful mean values
        cv_curve = np.zeros_like(mean_curve)
        cv_curve[mask] = (std_curve[mask] / mean_curve[mask]) * 100
        ax4.plot(time_points, cv_curve, linewidth=2, color=self.colors[3])
        ax4.set_xlabel('Time (hours)')
        ax4.set_ylabel('Coefficient of Variation (%)')
        ax4.set_title('Concentration CV vs Time Curve')
        ax4.grid(True, alpha=0.3)
        
        plt.suptitle('Population Variability Analysis', fontsize=16)
        plt.tight_layout()
        
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        
        return fig
    
    def plot_parameter_sensitivity(self, results: Dict[str, Any], save_path: Optional[str] = None) -> plt.Figure:
        """
        Plot parameter sensitivity analysis
        
        Parameters:
            results: Simulation results dictionary
            save_path: Save path
            
        Returns:
            matplotlib.figure.Figure: Chart object
        """
        # This is a simplified sensitivity analysis example
        fig, ax = plt.subplots(figsize=self.figsize)
        
        # Simulate some sensitivity data
        parameters = ['ka', 'Vd', 'CL', 'fu', 'pKa']
        sensitivity_values = np.random.randn(len(parameters)) * 0.5
        
        colors = ['green' if val > 0 else 'red' for val in sensitivity_values]
        ax.barh(parameters, sensitivity_values, color=colors, alpha=0.7)
        ax.set_xlabel('Sensitivity Coefficient')
        ax.set_title('Parameter Sensitivity Analysis')
        ax.grid(True, alpha=0.3)
        ax.axvline(x=0, color='black', linewidth=0.5)
        
        if save_path:
            fig.savefig(save_path, dpi=300, bbox_inches='tight')
        
        return fig

    def save_plots_batch(self, pbpk_results: Dict[str, Any], output_dir: str) -> Dict[str, str]:
        """
        Save all charts in batch
        
        Parameters:
            pbpk_results: PBPK modeling results dictionary
            output_dir: Output directory path
            
        Returns:
            dict: Saved file paths dictionary
        """
        import os
        os.makedirs(output_dir, exist_ok=True)
        
        saved_files = {}
        time_points = pbpk_results.get('time_points', [])
        plasma_conc = pbpk_results.get('plasma_concentration', [])
        tissue_concs = pbpk_results.get('tissue_concentrations', {})
        absorption_data = pbpk_results.get('absorption', {})
        elimination_data = pbpk_results.get('elimination', {})
        population_data = pbpk_results.get('population_data', {})
        
        try:
            # Plasma concentration curve
            if len(plasma_conc) > 0 and len(time_points) > 0:
                fig = self.plot_plasma_concentration(time_points, plasma_conc)
                path = os.path.join(output_dir, 'plasma_concentration.png')
                fig.savefig(path, dpi=300, bbox_inches='tight')
                saved_files['plasma_concentration'] = path
                plt.close(fig)
            
            # Tissue concentration curves
            if tissue_concs and len(time_points) > 0:
                fig = self.plot_tissue_concentrations(time_points, tissue_concs)
                path = os.path.join(output_dir, 'tissue_concentrations.png')
                fig.savefig(path, dpi=300, bbox_inches='tight')
                saved_files['tissue_concentrations'] = path
                plt.close(fig)
            
            # Absorption kinetics
            if absorption_data is not None and len(time_points) > 0:
                fig = self.plot_absorption_profile(time_points, absorption_data)
                path = os.path.join(output_dir, 'absorption_profile.png')
                fig.savefig(path, dpi=300, bbox_inches='tight')
                saved_files['absorption_profile'] = path
                plt.close(fig)
            
            # Elimination kinetics
            if elimination_data is not None and len(time_points) > 0:
                fig = self.plot_elimination_profile(time_points, elimination_data)
                path = os.path.join(output_dir, 'elimination_profile.png')
                fig.savefig(path, dpi=300, bbox_inches='tight')
                saved_files['elimination_profile'] = path
                plt.close(fig)
            
            # Population distribution
            if population_data is not None and len(population_data) > 0:
                fig = self.plot_population_distribution(population_data)
                path = os.path.join(output_dir, 'population_distribution.png')
                fig.savefig(path, dpi=300, bbox_inches='tight')
                saved_files['population_distribution'] = path
                plt.close(fig)
            
            # Summary dashboard
            fig = self.plot_pbpk_summary_dashboard(pbpk_results)
            path = os.path.join(output_dir, 'summary_dashboard.png')
            fig.savefig(path, dpi=300, bbox_inches='tight')
            saved_files['summary_dashboard'] = path
            plt.close(fig)
            
        except Exception as e:
            warnings.warn(f"Error occurred during chart saving: {str(e)}")
        
        return saved_files
    
    def create_animated_plasma_curve(self, all_simulations: List[Dict], output_path: str) -> str:
        """
        Create animated plasma concentration-time curve
        
        Parameters:
            all_simulations: List of all simulation results
            output_path: Output animation file path
            
        Returns:
            str: Animation file path
        """
        try:
            import matplotlib.animation as animation
            fig, ax = plt.subplots(figsize=self.figsize)
            ax.set_xlim(0, max([s['time_points'][-1] for s in all_simulations]))
            ax.set_ylim(0, max([max(s['plasma_concentration']) for s in all_simulations]) * 1.1)
            ax.set_xlabel('Time (hours)')
            ax.set_ylabel('Plasma Concentration (mg/L)')
            ax.set_title('Plasma Concentration-Time Curve Animation')
            ax.grid(True, alpha=0.3)
            line, = ax.plot([], [], lw=2)
            time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)
            
            def init():
                line.set_data([], [])
                time_text.set_text('')
                return line, time_text
            
            def animate(i):
                sim = all_simulations[i]
                line.set_data(sim['time_points'], sim['plasma_concentration'])
                time_text.set_text(f'Simulation {i+1}/{len(all_simulations)}')
                return line, time_text
            
            anim = animation.FuncAnimation(fig, animate, init_func=init,
                                         frames=len(all_simulations), interval=1000, blit=True)
            anim.save(output_path, writer='pillow', fps=1)
            plt.close(fig)
            return output_path
            
        except ImportError:
            warnings.warn("matplotlib.animation is required to create animations")
            return ""